import pandas as pd
import io # 用于创建内存中的CSV文件，方便演示




table1_path = "/sglang/outputs/analysis_performance_L40.csv"
table2_path = "/sglang/outputs/average_simple_time_L40.csv"
output_path = "/sglang/outputs/analysis_performance_with_acceleration_L40.csv"

# 实际读取文件的方式:
df1 = pd.read_csv(table1_path)
df2 = pd.read_csv(table2_path)

# --- 核心逻辑 ---

# 1. 将表2的step_time合并到表1中，基于batch_size
#    使用左合并 (left merge) 以保留表1中的所有行
merged_df = pd.merge(df1, df2, on="batch_size", how="left")

# 2. 计算 accelerate_rate
#    确保列名完全匹配你CSV文件中的列名
#    如果某batch_size在表2中找不到，对应的step_time会是NaN，计算结果也是NaN
merged_df['accelerate_rate'] = \
    (merged_df['nospec_step_time'] / merged_df['total_forward_batch_spec_time_ms']) * \
    merged_df['average_generate_tokens_per_step']

# --- 输出结果 ---

# 打印结果到控制台 (可选)
print("表1预览:")
print(df1.head())
print("\n表2预览:")
print(df2.head())
print("\n合并并计算后的表预览:")
print(merged_df.head())

# 将结果保存回CSV文件
# 你可以选择覆盖原始表1，或者保存到新文件
# output_path = table1_path # 覆盖原文件
merged_df.to_csv(output_path, index=False) # index=False 避免写入pandas的行索引

print(f"\n已将带有 'accelerate_rate' 列的表保存到: {output_path}")